#!/usr/bin/env python3
"""
DataPlexus 数据源管理权限初始化脚本
创建数据源管理相关的权限记录并分配给适当的角色
"""

import sys
import os
import asyncio
from datetime import datetime

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, and_
from app.core.database import get_async_session
from app.models.auth import Permission, Role, RolePermission
from app.core.logger import logger


class DataSourcePermissionInitializer:
    """数据源权限初始化器"""
    
    def __init__(self):
        self.permissions_data = [
            {
                'name': 'datasource.view',
                'display_name': '查看数据源',
                'description': '查看数据源列表和详情的权限',
                'resource': 'datasource',
                'action': 'view'
            },
            {
                'name': 'datasource.create',
                'display_name': '创建数据源',
                'description': '创建新数据源配置的权限',
                'resource': 'datasource',
                'action': 'create'
            },
            {
                'name': 'datasource.update',
                'display_name': '更新数据源',
                'description': '修改数据源配置的权限',
                'resource': 'datasource',
                'action': 'update'
            },
            {
                'name': 'datasource.delete',
                'display_name': '删除数据源',
                'description': '删除数据源配置的权限',
                'resource': 'datasource',
                'action': 'delete'
            },
            {
                'name': 'datasource.test',
                'display_name': '测试数据源连接',
                'description': '测试数据源连接状态的权限',
                'resource': 'datasource',
                'action': 'test'
            },
            {
                'name': 'datasource.import',
                'display_name': '导入数据源配置',
                'description': '批量导入数据源配置的权限',
                'resource': 'datasource',
                'action': 'import'
            },
            {
                'name': 'datasource.export',
                'display_name': '导出数据源配置',
                'description': '导出数据源配置的权限',
                'resource': 'datasource',
                'action': 'export'
            },
            {
                'name': 'datasource.manage',
                'display_name': '管理数据源',
                'description': '数据源管理的完整权限（包含所有操作）',
                'resource': 'datasource',
                'action': 'manage'
            }
        ]
        
        # 需要分配权限的角色
        self.role_permissions_mapping = {
            'admin': [
                'datasource.view', 'datasource.create', 'datasource.update', 
                'datasource.delete', 'datasource.test', 'datasource.import', 
                'datasource.export', 'datasource.manage'
            ],
            'data_admin': [
                'datasource.view', 'datasource.create', 'datasource.update', 
                'datasource.test', 'datasource.import', 'datasource.export'
            ],
            'developer': [
                'datasource.view', 'datasource.test'
            ],
            'analyst': [
                'datasource.view'
            ]
        }
    
    async def create_permissions(self, db: AsyncSession) -> dict:
        """创建数据源管理权限"""
        created_count = 0
        updated_count = 0
        skipped_count = 0
        
        try:
            for perm_data in self.permissions_data:
                # 检查权限是否已存在
                stmt = select(Permission).where(Permission.name == perm_data['name'])
                result = await db.execute(stmt)
                existing_permission = result.scalar_one_or_none()
                
                if existing_permission:
                    # 更新现有权限
                    existing_permission.display_name = perm_data['display_name']
                    existing_permission.description = perm_data['description']
                    existing_permission.resource = perm_data['resource']
                    existing_permission.action = perm_data['action']
                    existing_permission.updated_at = datetime.utcnow()
                    updated_count += 1
                    logger.info(f"更新权限: {perm_data['name']}")
                else:
                    # 创建新权限
                    permission = Permission(
                        name=perm_data['name'],
                        display_name=perm_data['display_name'],
                        description=perm_data['description'],
                        resource=perm_data['resource'],
                        action=perm_data['action'],
                        is_system=True,
                        created_at=datetime.utcnow(),
                        updated_at=datetime.utcnow()
                    )
                    db.add(permission)
                    created_count += 1
                    logger.info(f"创建权限: {perm_data['name']}")
            
            await db.commit()
            
            return {
                'created': created_count,
                'updated': updated_count,
                'skipped': skipped_count,
                'total': len(self.permissions_data)
            }
            
        except Exception as e:
            await db.rollback()
            logger.error(f"创建权限失败: {str(e)}")
            raise
    
    async def assign_permissions_to_roles(self, db: AsyncSession) -> dict:
        """将权限分配给角色"""
        assigned_count = 0
        skipped_count = 0
        
        try:
            for role_name, permission_names in self.role_permissions_mapping.items():
                # 获取角色
                role_stmt = select(Role).where(Role.name == role_name)
                role_result = await db.execute(role_stmt)
                role = role_result.scalar_one_or_none()
                
                if not role:
                    logger.warning(f"角色不存在: {role_name}")
                    continue
                
                for permission_name in permission_names:
                    # 获取权限
                    perm_stmt = select(Permission).where(Permission.name == permission_name)
                    perm_result = await db.execute(perm_stmt)
                    permission = perm_result.scalar_one_or_none()
                    
                    if not permission:
                        logger.warning(f"权限不存在: {permission_name}")
                        continue
                    
                    # 检查角色权限关联是否已存在
                    rp_stmt = select(RolePermission).where(
                        and_(
                            RolePermission.role_id == role.id,
                            RolePermission.permission_id == permission.id
                        )
                    )
                    rp_result = await db.execute(rp_stmt)
                    existing_rp = rp_result.scalar_one_or_none()
                    
                    if existing_rp:
                        skipped_count += 1
                        continue
                    
                    # 创建角色权限关联
                    role_permission = RolePermission(
                        role_id=role.id,
                        permission_id=permission.id,
                        is_active=True,
                        created_at=datetime.utcnow()
                    )
                    db.add(role_permission)
                    assigned_count += 1
                    logger.info(f"为角色 {role_name} 分配权限: {permission_name}")
            
            await db.commit()
            
            return {
                'assigned': assigned_count,
                'skipped': skipped_count
            }
            
        except Exception as e:
            await db.rollback()
            logger.error(f"分配权限失败: {str(e)}")
            raise
    
    async def verify_permissions(self, db: AsyncSession) -> dict:
        """验证权限创建结果"""
        try:
            # 统计数据源相关权限
            perm_stmt = select(Permission).where(Permission.resource == 'datasource')
            perm_result = await db.execute(perm_stmt)
            permissions = perm_result.scalars().all()
            
            # 统计角色权限关联
            rp_stmt = select(RolePermission).join(Permission).where(
                Permission.resource == 'datasource'
            )
            rp_result = await db.execute(rp_stmt)
            role_permissions = rp_result.scalars().all()
            
            verification_result = {
                'permissions_count': len(permissions),
                'role_permissions_count': len(role_permissions),
                'permissions': [
                    {
                        'name': p.name,
                        'display_name': p.display_name,
                        'resource': p.resource,
                        'action': p.action
                    } for p in permissions
                ]
            }
            
            return verification_result
            
        except Exception as e:
            logger.error(f"验证权限失败: {str(e)}")
            raise
    
    async def run(self):
        """运行权限初始化"""
        logger.info("开始初始化数据源管理权限...")
        
        try:
            async with get_async_session() as db:
                # 创建权限
                logger.info("创建数据源管理权限...")
                perm_result = await self.create_permissions(db)
                logger.info(f"权限创建完成: {perm_result}")
                
                # 分配权限给角色
                logger.info("分配权限给角色...")
                assign_result = await self.assign_permissions_to_roles(db)
                logger.info(f"权限分配完成: {assign_result}")
                
                # 验证结果
                logger.info("验证权限创建结果...")
                verify_result = await self.verify_permissions(db)
                logger.info(f"验证结果: {verify_result}")
                
                logger.info("数据源管理权限初始化完成！")
                
                return {
                    'permissions': perm_result,
                    'assignments': assign_result,
                    'verification': verify_result
                }
                
        except Exception as e:
            logger.error(f"权限初始化失败: {str(e)}")
            raise


async def main():
    """主函数"""
    try:
        initializer = DataSourcePermissionInitializer()
        result = await initializer.run()
        
        print("\n" + "="*60)
        print("DataPlexus 数据源管理权限初始化完成")
        print("="*60)
        print(f"创建权限: {result['permissions']['created']}")
        print(f"更新权限: {result['permissions']['updated']}")
        print(f"分配权限: {result['assignments']['assigned']}")
        print(f"跳过权限: {result['assignments']['skipped']}")
        print(f"总权限数: {result['verification']['permissions_count']}")
        print(f"角色权限关联数: {result['verification']['role_permissions_count']}")
        print("="*60)
        
        return True
        
    except Exception as e:
        print(f"\n❌ 权限初始化失败: {str(e)}")
        return False


if __name__ == "__main__":
    success = asyncio.run(main())
    sys.exit(0 if success else 1)
